"""
Text2Vec Base Chinese 模型加载器
支持昇腾910B NPU加速
"""
import os
import torch
import logging
from typing import Optional
from sentence_transformers import SentenceTransformer

# 条件导入torch_npu
try:
    import torch_npu
    NPU_AVAILABLE = True
except ImportError:
    NPU_AVAILABLE = False

logger = logging.getLogger(__name__)


class Text2VecChineseModelLoader:
    """Text2Vec Chinese模型加载器"""
    
    def __init__(self, model_id: str = "shibing624/text2vec-base-chinese", cache_dir: Optional[str] = None, local_model_path: Optional[str] = None):
        self.model_id = model_id
        self.cache_dir = cache_dir or os.getenv("MODEL_CACHE_DIR", "/app/models")
        self.local_model_path = local_model_path or os.getenv("LOCAL_MODEL_PATH", None)
        self.device = None
        self.model = None
        
    def _get_device(self):
        """获取计算设备"""
        if NPU_AVAILABLE and hasattr(torch, 'npu') and torch.npu.is_available():
            device = "npu:0"
            logger.info("使用昇腾NPU加速")
        elif torch.cuda.is_available():
            device = "cuda"
            logger.info("使用CUDA加速")
        else:
            device = "cpu"
            logger.info("使用CPU")
        return device
    
    def load_model(self):
        """加载模型"""
        try:
            self.device = self._get_device()
            
            # 检查是否有本地模型路径
            if self.local_model_path and os.path.exists(self.local_model_path):
                logger.info(f"从本地路径加载模型: {self.local_model_path}")
                model_path = self.local_model_path
            else:
                logger.info(f"从HuggingFace加载模型: {self.model_id}")
                model_path = self.model_id
            
            # 使用SentenceTransformer加载模型
            self.model = SentenceTransformer(
                model_path,
                cache_folder=self.cache_dir,
                device=self.device
            )
            
            logger.info("模型加载成功")
            return True
            
        except Exception as e:
            logger.error(f"模型加载失败: {str(e)}")
            raise
    
    def get_model(self):
        """获取模型"""
        if self.model is None:
            raise RuntimeError("模型未加载，请先调用load_model()")
        return self.model
    
    def get_device(self):
        """获取设备"""
        return self.device


